from torch.utils.data import DataLoader

from ARM_TGNCF.Parse_ARM_TGNCF import ARM_args
from loaddata.Ratings_Dataset import Ratings_Dataset


def get_UI_type():
    if ARM_args.UI_type == 0:
        uu = False
        ii = False
        method_inf = "no auxiliary  graph"

    elif ARM_args.UI_type == 1:
        uu = True
        ii = False
        method_inf = "only user,and the links is {}".format(ARM_args.u_top_k)

    elif ARM_args.UI_type == 2:
        uu = False
        ii = True
        method_inf = "only item,and the links is {}".format(ARM_args.u_top_k)

    else:
        uu = True
        ii = True
        method_inf = "user and item,and the links is {} and {}".format(ARM_args.u_top_k, ARM_args.i_top_k)

    return uu, ii, method_inf


def get_DataLoader(train_event, test_event):
    user_num = max(train_event['userId'].max(), test_event['userId'].max()) + 1
    item_num = max(train_event['itemId'].max(), test_event['itemId'].max()) + 1
    print('user_num:', user_num, 'item_num:', item_num)

    train = Ratings_Dataset(train_event)
    test = Ratings_Dataset(test_event)

    dl_train = DataLoader(train, batch_size=ARM_args.batch_size, shuffle=not ARM_args.events_is_ascend_time,
                          pin_memory=True)
    dl_test = DataLoader(test, batch_size=len(test))

    return user_num, item_num, dl_train, dl_test


def print_args_information(method_inf, user_num, item_num, ratings_num):
    info = 'dataset:{}'.format(ARM_args.dataset) + '\n'

    train_size = int(ratings_num * ARM_args.train)
    test_size = ratings_num - train_size
    info += 'events num:{},{} for train,{} for test'.format(ratings_num, train_size, test_size) + '\n'
    info += 'user_num: {} item_num: {}'.format(user_num, item_num)

    info += 'item Auxiliary method：{}'.format(ARM_args.ii_method)
    if ARM_args.ARM_sequential:
        info += 'item time '
    if ARM_args.is_in_train:
        info += 'dy'
    if ARM_args.events_is_ascend_time:
        info += 'event time'
    info += '\n'
    info += method_inf + '\n'

    info += "epoch={},batch_size={},lr={},seed={},embed_size={},decay={}".format(
        ARM_args.epoch, ARM_args.batch_size, ARM_args.lr, ARM_args.seed, ARM_args.embed_size, ARM_args.decay) + '\n'
    print(info, '\n')
    return info
